#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project : RAG-demo
# @File    : sqlite_util.py
# @IDE     : PyCharm
# @Author  :ZH
# @Time    : 2025/7/2 10:20
import sqlite3


class SQLiteDB:
    def __init__(self, db_path: str):
        self.db_path = db_path
        self.connection = None
        self.cursor = None

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def connect(self):
        self.connection = sqlite3.connect(self.db_path)
        self.cursor = self.connection.cursor()

    def close(self):
        if self.cursor:
            self.cursor.close()
        if self.connection:
            self.connection.close()

    def table_exists(self, table_name: str) -> bool:
        """
        判断指定的表是否已经存在
        :param table_name: 表名
        :return: True 如果存在，否则 False
        """
        sql = "SELECT name FROM sqlite_master WHERE type='table' AND name=?"
        result = self.execute(sql, (table_name,))
        return len(result) > 0

    def execute(self, sql: str, params=None):
        if params is None:
            params = ()
        try:
            self.cursor.execute(sql, params)
            if sql.strip().lower().startswith("select"):
                return self.cursor.fetchall()
            else:
                self.connection.commit()
                return self.cursor.rowcount
        except sqlite3.Error as e:
            self.connection.rollback()
            raise e

    def execute_many(self, sql: str, params_list):
        try:
            self.cursor.executemany(sql, params_list)
            self.connection.commit()
            return self.cursor.rowcount
        except sqlite3.Error as e:
            self.connection.rollback()
            raise e

    def create_table(self, table_sql: str):
        self.execute(table_sql)

    def fetch_one(self, sql: str, params=None):
        if params is None:
            params = ()
        self.cursor.execute(sql, params)
        return self.cursor.fetchone()

    def fetch_all(self, sql: str, params=None):
        if params is None:
            params = ()
        self.cursor.execute(sql, params)
        return self.cursor.fetchall()

    def insert(self, table: str, data: dict):
        columns = ', '.join(data.keys())
        placeholders = ', '.join('?' * len(data))
        values = tuple(data.values())
        sql = f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"
        return self.execute(sql, values)

    def update(self, table: str, data: dict, where_clause: str, where_params: tuple):
        set_clause = ', '.join([f"{key} = ?" for key in data.keys()])
        values = tuple(data.values()) + where_params
        sql = f"UPDATE {table} SET {set_clause} WHERE {where_clause}"
        return self.execute(sql, values)

    def delete(self, table: str, where_clause: str, where_params: tuple):
        sql = f"DELETE FROM {table} WHERE {where_clause}"
        return self.execute(sql, where_params)
